import math, os
import PyPDF2
from reportlab.lib import pagesizes
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont

pdfmetrics.registerFont(TTFont("STSongti", "Songti.ttc"))


class WaterMarkUtils:
    content: str
    width: int
    height: int
    text_stroke_color_rgb = (0, 0, 0)
    text_fill_color_rgb = (0, 0, 0)
    text_fill_alpha = 0.1
    page_size: tuple
    cv: canvas.Canvas
    n = 3

    def __init__(self, content, page_size=pagesizes.A4):
        self.content = content
        self.page_size = page_size
        self.width = self.page_size[0]
        self.height = self.page_size[1] / self.n
        self.draw_pdf()

    @staticmethod
    def lenb(str_obj):
        str_len = 0
        for s in str_obj:
            str_len += 1 if ord(s) < 0x4E00 else 2
        return str_len

    @property
    def auto_angle(self):
        return math.atan(self.height / self.width) / 2 / math.pi * 360

    @property
    def content_length(self):
        return self.lenb(self.content) / 2

    @property
    def diagonal_length(self):
        return (self.width**2 + self.height**2) ** 0.5

    @property
    def font_size(self):
        size = self.diagonal_length / (
            self.content_length + self.width / self.height + 1
        )
        return size

    @property
    def draw_x(self):
        return ((self.width / self.height + 1) * self.font_size) / 2

    @property
    def draw_y(self):
        return -self.font_size / 2

    @property
    def move_y(self):
        return self.width / self.diagonal_length * self.height

    @property
    def move_x(self):
        return self.height**2 / self.diagonal_length

    def draw_pdf(self):
        self.cv = canvas.Canvas(f"{self.content}.pdf", pagesize=self.page_size)
        self.cv.setFont("STSongti", self.font_size)
        self.cv.setStrokeColorRGB(*self.text_stroke_color_rgb)
        self.cv.setFillColorRGB(*self.text_fill_color_rgb)
        self.cv.setFillAlpha(self.text_fill_alpha)
        self.cv.rotate(self.auto_angle)
        for i in range(self.n):
            self.cv.drawString(
                self.draw_x + i * self.move_x,
                self.draw_y + i * self.move_y,
                self.content,
            )

        # self.cv.drawString(self.draw_x, self.draw_y, self.content)

    def save(self):
        self.cv.save()
        return f"{self.content}.pdf"


def get_page_size(page: PyPDF2._page.PageObject):
    page_size = page.mediabox
    return float(page_size[2]), float(page_size[3])


def add_watermark(content, source_path: str, out_path: str = None):
    target_file_name = out_path or source_path.split(".")[0] + "_wm.pdf"
    wm_file = f"{content}.pdf"
    wm = {}
    with open(source_path, "rb") as source_file:
        source_reader = PyPDF2.PdfReader(source_file)
        target_writer = PyPDF2.PdfWriter()
        for page in source_reader.pages:
            page_size = page.mediabox
            page_size_str = str(page_size)
            if page_size_str not in wm.keys():
                if not os.path.exists(wm_file):
                    wm_file = WaterMarkUtils(content, (get_page_size(page))).save()
                wm_reader = PyPDF2.PdfReader(wm_file)
                wm[page_size_str] = wm_reader.pages[0]
            page.merge_page(wm[page_size_str])
            target_writer.add_page(page)
        with open(target_file_name, "wb") as write_file:
            target_writer.write(write_file)





if __name__ == "__main__":
    add_watermark(
        "天衡会计师事务所（特殊普通合伙）苏州安信分所", "4264.pdf", "4264.pdf"
    )
    # WaterMarkUtils("草稿2022年5月26日19:11:08").save()
